import cv2
import subprocess
import numpy as np
from flask import Flask, request, jsonify
from threading import Thread, Lock

import time

from tools.reco.recognition_camera import predict_expression
from. import reco_app


# 全局状态管理（使用线程锁保证安全）
stream_control = {
    "is_running": False,
    "ffmpeg_input": None,
    "ffmpeg_output": None,
    "frame_size": 640 * 480 * 3,
    "lock": Lock()
}

def process_stream():
    with stream_control['lock']:
        if not stream_control['is_running']:
            return

        # 初始化FFmpeg进程
        try:
            # 拉流进程
            stream_control['ffmpeg_input'] = subprocess.Popen([
                'ffmpeg',
                '-i', 'rtmp://192.168.31.105:1935/live_in/1',
                '-f', 'image2pipe',
                '-pix_fmt', 'bgr24',    # 输入为 BGR 格式
                '-vcodec', 'rawvideo',
                '-s', '640x480',        # 分辨率 640x480
                '-r', '30',
                '-'
            ], stdout=subprocess.PIPE, stderr=subprocess.PIPE)

            # 推流进程
            stream_control['ffmpeg_output'] = subprocess.Popen([
                'ffmpeg',
                '-y',
                '-f', 'rawvideo',
                '-vcodec', 'rawvideo',
                '-pix_fmt', 'bgr24',    # 输出保持 BGR 格式
                '-s', '640x480',
                '-r', '30',
                '-i', '-',
                '-c:v', 'libx264',
                '-preset', 'ultrafast',
                '-tune', 'zerolatency',
                '-f', 'flv',
                'rtmp://192.168.31.105:1935/live_out/1'
            ], stdin=subprocess.PIPE, stderr=subprocess.PIPE)
        except Exception as e:
            print("FFmpeg启动失败:", str(e))
            stream_control['is_running'] = False
            return

    # 流处理主循环
    while True:
        with stream_control['lock']:
            if not stream_control['is_running']:
                break

            # 检查进程状态
            if stream_control['ffmpeg_input'].poll() is not None:
                error = stream_control['ffmpeg_input'].stderr.read()
                print("拉流进程异常退出:", error)
                break
            if stream_control['ffmpeg_output'].poll() is not None:
                error = stream_control['ffmpeg_output'].stderr.read()
                print("推流进程异常退出:", error)
                break

        # 读取帧数据
        try:
            raw_frame = stream_control['ffmpeg_input'].stdout.read(640 * 480 * 3)
            if not raw_frame:
                print("流数据读取结束")
                break

            # 转换为 numpy 数组 (height, width, channels)
            frame = np.frombuffer(raw_frame, dtype=np.uint8).reshape(480, 640, 3)

            # 处理帧（确保返回 BGR 格式）
            processed_frame = predict_expression(frame)

            # 强制调整分辨率（如果处理过程中修改了尺寸）
            if processed_frame.shape[:2] != (480, 640):
                processed_frame = cv2.resize(processed_frame, (640, 480))

            # 推流（直接写入 BGR 数据）
            stream_control['ffmpeg_output'].stdin.write(processed_frame.tobytes())
            stream_control['ffmpeg_output'].stdin.flush()

        except Exception as e:
            print("帧处理异常:", str(e))
            break

    # 清理资源
    with stream_control['lock']:
        if stream_control['ffmpeg_input'] and stream_control['ffmpeg_input'].poll() is None:
            stream_control['ffmpeg_input'].terminate()
        if stream_control['ffmpeg_output'] and stream_control['ffmpeg_output'].poll() is None:
            stream_control['ffmpeg_output'].terminate()
        stream_control['is_running'] = False

@reco_app.route('/start', methods=['POST'])
def start_stream():
    print("启动流")
    with stream_control['lock']:
        if stream_control['is_running']:
            return jsonify({"status": "error", "message": "流已在运行"})

        stream_control['is_running'] = True
        # 启动流处理线程
        Thread(target=process_stream).start()
        return jsonify({"status": "success", "message": "流已启动"})


@reco_app.route('/stop', methods=['POST'])
def stop_stream():
    with stream_control['lock']:
        if not stream_control['is_running']:
            return jsonify({"status": "error", "message": "流未运行"})

        # 设置停止标志
        stream_control['is_running'] = False
        return jsonify({"status": "success", "message": "停止指令已发送"})
